from dataset import Data, Data2
from sklearn import tree
from sklearn.metrics import mean_absolute_percentage_error, mean_absolute_error, mean_squared_error
from visual import plot, dfplot, plot2
import pandas as pd
import numpy as np

class Trainer:
    def __init__(self, args):
        self.args = args

        self._init_data()
        
    def _init_data(self):
        data = Data2(self.args.data_path)
        self.train_date = data.train_date
        self.test_date = data.test_date
        self.train_X = data.train_X
        self.test_X = data.test_X
        self.train_y = data.train_y
        self.test_y = data.test_y



    def train_dt(self):
        self.model = tree.DecisionTreeRegressor()
        self.model.fit(self.train_X, self.train_y)


    def test(self):
        print(self.test_X.shape)
        pred = self.model.predict(self.test_X)
        mae = mean_absolute_percentage_error(self.test_y, pred)
        print("mae_loss: %.4f" % (mae))

        df = pd.DataFrame()
        df['date'] = self.test_date
        df['y'] = self.test_y
        df['pred'] = pred
        new_df = df.groupby('date').mean()
        # print(new_df)
        plot(new_df.index.tolist(), [new_df.y, new_df.pred], ['真实值', '预测值'], 'date', 'value', './results/')
        
    def test2(self):
        """
        散点图
        """
        self.test_X = np.concatenate([self.train_X, self.test_X], axis=0)
        self.test_y = np.concatenate([self.train_y, self.test_y])
        pred = self.model.predict(self.test_X)
        # mse = mean_squared_error(self.test_y, pred)
        # mae = mean_absolute_error(self.test_y, pred)
        mse = np.sqrt(np.mean((self.test_y - pred)**2))      # 均方根误差
        mae = np.mean(np.abs(self.test_y - pred))            # 平均绝对值误差
        print("均方根误差: %.4f" % (mse))  
        print("平均绝对误差: %.4f" % mae)
        # print(pred.shape, self.test_y.shape)
        # df = pd.DataFrame()
        # df['date'] = self.test_date
        # df['y'] = self.test_y
        # df['pred'] = pred
        plot2(self.test_y, pred, xlabel="Actual Value", ylabel="Estimate", save_path="./results")


